# import numpy as np
import itertools


def _compare_float(a, b):
    precision = 1e-10
    if abs(a - b) <= precision:
        return 1
    return 0


def compare_float(a, b):
    precision = 1e-10
    if abs(a - b) <= precision:
        return True
    return False


def judge_complex(c):
    index = [0, 1, -1]
    _sum = sum(
        itertools.starmap(
            lambda x, y: _compare_float(c.real, x) * _compare_float(c.imag, y),
            itertools.product(index, index),
        )
    )
    return 1 - _sum
    # if (compare_float(c.real, 0) and compare_float(c.imag, 0)):
    #     return 0
    # elif (compare_float(c.real, 1) and compare_float(c.imag, 0)):
    #     return 0
    # elif (compare_float(c.real, -1) and compare_float(c.imag, 0)):
    #     return 0
    # elif (compare_float(c.real, 0) and compare_float(c.imag, 1)):
    #     return 0
    # elif (compare_float(c.real, 0) and compare_float(c.imag, -1)):
    #     return 0
    # elif (compare_float(c.real, 1) and compare_float(c.imag, 1)):
    #     return 0
    # elif (compare_float(c.real, 1) and compare_float(c.imag, -1)):
    #     return 0
    # elif (compare_float(c.real, -1) and compare_float(c.imag, 1)):
    #     return 0
    # elif (compare_float(c.real, -1) and compare_float(c.imag, -1)):
    #     return 0
    # else:
    #     return 1


def compute_complex(A1, A2):
    complex_degree = 0
    matrix_result = []
    for i in range(len(A1)):  # A1矩阵的行
        row = []
        for j in range(len(A2[0])):  # A2矩阵的列
            sum = 0
            for k in range(len(A2)):
                if judge_complex(A1[i][k]) == 1 and judge_complex(A2[k][j]) == 1:
                    complex_degree += 1
                sum += A1[i][k] * A2[k][j]
            row.append(sum)
        matrix_result.append(row)
    return matrix_result, complex_degree


def _compute_complex(A1, A2):
    complex_degree = 0
    # for i in range(len(A1)):  # A1矩阵的行
    complex_degree = sum(
        map(
            lambda A1_i: sum(
                map(
                    lambda j: sum(
                        itertools.starmap(
                            lambda x, y: judge_complex(x) * judge_complex(y[j]),
                            zip(A1_i, A2),
                        )
                    ),
                    range(len(A2)),
                )
            ),
            A1,
        )
    )

    # for j in range(len(A2)):  # A2矩阵的列
    #     complex_degree += sum(itertools.starmap(lambda x,
    #                           y: judge_complex(x)*judge_complex(y[j]), zip(A1[i], A2)))
    return complex_degree


A = [[2, 2 + 4j, -1], [1 + 2j, 0, -1j], [1 + 8j, 0, 1j]]
B = [[0, 1, 2 + 1j], [2 + 4j, 2 - 4j, -4j], [1 + 1j, 0, 1j]]
# A = [[1, 2+4j], [1+2j, 0]]
# B = [[0, 1], [2+4j, 2-4j]]
print(_compute_complex(A, B))
